#!/usr/bin/env python
# coding: utf-8

# if in rstudio: reticulate::use_condaenv('scvelo-colorbar')

# ======= Differential Kinetics =====
# sample with multiple lineages and processes,
# where genes are likely to show different kinetic regimes across subpopulations
import scvelo as scv

# prepare the data ----------
adata = scv.read('/home/gjsx/learn/scvelo/data/DentateGyrus/10X43_1.h5ad', cache=True)

scv.pp.filter_and_normalize(adata, min_shared_counts=30, n_top_genes=2000)
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)

# basic velocity estimation --------
scv.tl.velocity(adata)
scv.tl.velocity_graph(adata)

scv.pl.velocity_embedding_stream(adata, basis='umap')

# Differential Kinetic Test ----------
# choose 4 genes to test for each clusters, default will use all velocity genes
var_names = ['Tmsb10', 'Fam155a', 'Hn1', 'Rpl6']
scv.tl.differential_kinetic_test(adata, var_names=var_names, groupby='clusters')

# best fit clusters of 4 genes?
scv.get_df(adata[:, var_names], ['fit_diff_kinetics', 'fit_pval_kinetics'], precision=2)

# plot
kwargs = dict(linewidth=2, add_linfit=True, frameon=False)
scv.pl.scatter(adata, basis=var_names, add_outline='fit_diff_kinetics', **kwargs)
# In Tmsb10, for instance, Endothelial display a kinetic behaviour (outlined, with the black line fitted through),
# that cannot be well explained by the overall dynamics (purple curve).

diff_clusters=list(adata[:, var_names].var['fit_diff_kinetics'])
scv.pl.scatter(adata, legend_loc='right', size=60, title='diff kinetics',
               add_outline=diff_clusters, outline_width=(.8, .2))
               
# Testing top-likelihood genes ------
# Screening through the top-likelihood genes in dynamic modeling
scv.tl.recover_dynamics(adata, n_jobs=16)

# find best fit clusters for 100 best overall fit genes
top_genes = adata.var['fit_likelihood'].sort_values(ascending=False).index[:100]
scv.tl.differential_kinetic_test(adata, var_names=top_genes, groupby='clusters')

scv.pl.scatter(adata, basis=top_genes[:15], ncols=5, add_outline='fit_diff_kinetics', **kwargs)

# Recompute velocities ----------
# velocities can be recomputed leveraging the information of multiple competing kinetic
# diff_kinetics param is not in the helper?
scv.tl.velocity(adata, diff_kinetics=True)
scv.tl.velocity_graph(adata)

scv.pl.velocity_embedding_stream(adata)
